"""
https://pypi.org/project/streamlit-chatbox/
https://pypi.org/project/streamlit-chatbox/1.1.11/
https://discuss.streamlit.io/t/new-component-chatbox/44151
"""

import streamlit as st
import streamlit_chatbox as stcb
import time
import simplejson as json
from PyCmpltrtok.common import sep

llm = stcb.FakeLLM()
chat_box = stcb.ChatBox()
sep('1')

with st.sidebar:
    sep('2')
    st.subheader('Start to chat with streamlit')
    streaming = st.checkbox('Streaming', True)
    in_expander = st.checkbox('Show message in expander', True)
    show_history = st.checkbox('Show history', False)
    
    st.divider()
    
    btns = st.container()
    
    file = st.file_uploader('Chat history JSON', type=['json'])
    
    if st.button('Load JSON') and file:
        data = json.load(file)
        chat_box.from_dict(data)
        
chat_box.init_session()
chat_box.output_messages()
    
if query := st.chat_input('Input your question here'):
    sep('3')
    chat_box.user_say(query)
    
    if streaming:
        generator = llm.chat_stream(query)
        elements = chat_box.ai_say([
            stcb.Markdown('Thinking', in_expander=in_expander, expanded=True, title="Answer"),
            stcb.Markdown("", in_expander=in_expander, title="references"),
        ])
        time.sleep(0.2)
        text = ''
        first = 1
        for x, docs in generator:
            if first:
                first = 0
                chat_box.update_msg('\n\n'.join(docs), element_index=1, streaming=False, state="complete", expanded=True)        
            text += x
            chat_box.update_msg(text, element_index=0, streaming=True)
            
        chat_box.update_msg(text, element_index=0, streaming=False, state="complete")
        
    else:
        text, docs = llm.chat(query)
        elements = chat_box.ai_say([
            stcb.Markdown(text, in_expander=in_expander, expanded=True, title="Answer", state='complete'),
            stcb.Markdown('\n\n'.join(docs), in_expander=in_expander, title="references", state='complete'),
        ])
sep('4')
        
cols = st.columns(2)
if cols[0].button('Show me the multimedia'):
    chat_box.ai_say(stcb.Image('https://tse4-mm.cn.bing.net/th/id/OIP-C.cy76ifbr2oQPMEs2H82D-QHaEv?w=284&h=181&c=7&r=0&o=5&dpr=1.5&pid=1.7'))
    time.sleep(0.5)
    chat_box.ai_say(stcb.Video('https://download.samplelib.com/mp4/sample-5s.mp4'))
    time.sleep(0.5)
    chat_box.ai_say(stcb.Audio('https://download.samplelib.com/mp4/sample-5s.mp4'))
    
if cols[1].button('Run agent'):
    chat_box.user_say('Run agent!')
    agent = stcb.FakeAgent()
    text = ''
    
    chat_box.ai_say()  # placeholder
    for d in agent.run_stream():
        if d['type'] == 'complete':
            chat_box.update_msg(expanded=False, state='complete')
            chat_box.insert_msg(d['llm_output'])
            break
        
        if d['status'] == 1:
            chat_box.update_msg(expanded=False, state="complete")
            text = ""
            chat_box.insert_msg(stcb.Markdown(text, title=d["text"], in_expander=True, expanded=True))
        elif d['status'] == 2:
            text += d['llm_output']
            chat_box.update_msg(text, streaming=True)
        else:
            chat_box.update_msg(text, streaming=False)
            
btns.download_button(
    "Export Markdown",
    "".join(chat_box.export2md()),
    file_name=f"chat_history.md",
    mime="text/markdown",
)

btns.download_button(
    "Export Json",
    chat_box.to_json(),
    file_name="chat_history.json",
    mime="text/json",
)

if btns.button("clear history"):
    chat_box.init_session(clear=True)
    st.experimental_rerun()
    
if show_history:
    st.write(chat_box.history)
    